#include "multiple.cuh"
#include <stdio.h>

// Matrix multiplication - Host code
// Matrix dimensions are assumed to be multiples of BLOCK_SIZE
void CUDAMatMul(const _PlainMatrix A, const _PlainMatrix B, _PlainMatrix C) {
  if (A.width != B.height) {
    fprintf(stderr, "Error: A.width (%d) must equal B.height (%d)\n", A.width,
            B.height);
    return;
  }
  // Load A and B to device memory
  _PlainMatrix d_A;
  d_A.width = A.width;
  d_A.height = A.height;
  size_t size = A.width * A.height * sizeof(float);
  cudaMalloc(&d_A.elements, size);
  cudaMemcpy(d_A.elements, A.elements, size, cudaMemcpyHostToDevice);
  _PlainMatrix d_B;
  d_B.width = B.width;
  d_B.height = B.height;
  size = B.width * B.height * sizeof(float);
  cudaMalloc(&d_B.elements, size);
  cudaMemcpy(d_B.elements, B.elements, size, cudaMemcpyHostToDevice);

  // Allocate C in device memory
  _PlainMatrix d_C;
  d_C.width = C.width;
  d_C.height = C.height;
  size = C.width * C.height * sizeof(float);
  cudaMalloc(&d_C.elements, size);

  // Invoke kernel
  dim3 dimBlock(BLOCK_SIZE, BLOCK_SIZE);
  //   dim3 dimGrid(B.width / dimBlock.x, A.height / dimBlock.y);
  dim3 dimGrid((C.width + dimBlock.x - 1) / dimBlock.x,
               (C.height + dimBlock.y - 1) / dimBlock.y);
  MatMulKernel<<<dimGrid, dimBlock>>>(d_A, d_B, d_C);

  // Read C from device memory
  cudaMemcpy(C.elements, d_C.elements, size, cudaMemcpyDeviceToHost);

  // Free device memory
  cudaFree(d_A.elements);
  cudaFree(d_B.elements);
  cudaFree(d_C.elements);
}

// Matrix multiplication kernel called by CUDAMatMul()
__global__ void MatMulKernel(_PlainMatrix A, _PlainMatrix B, _PlainMatrix C) {
  int row = blockIdx.y * blockDim.y + threadIdx.y;
  int col = blockIdx.x * blockDim.x + threadIdx.x;

  // 提前退出非法线程
  if (row >= C.height || col >= C.width)
    return;

  float Cvalue = 0.0f;
  for (int e = 0; e < A.width; ++e)
    Cvalue += A.elements[row * A.width + e] * B.elements[e * B.width + col];

  C.elements[row * C.width + col] = Cvalue;
}